import numpy as np
from matplotlib import pyplot as plt
# 注意：确保已安装 'pillow' 库以支持图像加载
from sklearn.datasets import load_sample_image
from sklearn.cluster import MiniBatchKMeans

# 加载 sklearn 数据集中的示例图片 "china.jpg"
china = load_sample_image("china.jpg")

# 显示原始图片
ax = plt.axes(xticks=[], yticks=[])  # 去除坐标轴刻度以提升可视化效果
ax.imshow(china)
plt.title("Original Image")
plt.show()

# 查看图片的形状
print("Image shape:", china.shape)  # 形状：(427, 640, 3)，代表高度、宽度和RGB通道

# 将图像数据归一化到 0 到 1 的范围
data = china / 255.0
# 将图像数据从 (427, 640, 3) 重塑为 (427*640, 3)，以便聚类
data = data.reshape(-1, 3)  # 自动计算第一个维度
print("Reshaped data shape:", data.shape)  # 形状：(273280, 3)


# 定义一个函数用于在 RGB 色彩空间中可视化像素分布
def plot_pixels(data, title, colors=None, N=10000):
    """
    在 RGB 色彩空间中绘制像素分布图

    参数:
        data (ndarray): 图像的 RGB 数据，形状为 (N, 3)
        title (str): 绘图标题
        colors (ndarray, optional): 用于绘制点的颜色。默认为 None（使用数据本身）
        N (int, optional): 随机采样的像素数量。默认值为 10,000
    """
    if colors is None:
        colors = data

    # 随机选择一部分像素点
    rng = np.random.RandomState(0)
    i = rng.permutation(data.shape[0])[:N]
    colors = colors[i]
    R, G, B = data[i].T  # 转置以分别提取 R、G、B 通道

    # 创建散点图以展示颜色分布
    fig, ax = plt.subplots(1, 2, figsize=(16, 6))
    ax[0].scatter(R, G, color=colors, marker='.')
    ax[0].set(xlabel='Red', ylabel='Green', xlim=(0, 1), ylim=(0, 1))
    ax[1].scatter(R, B, color=colors, marker='.')
    ax[1].set(xlabel='Red', ylabel='Blue', xlim=(0, 1), ylim=(0, 1))

    fig.suptitle(title, size=20)
    plt.show()


# 可视化原始图像的颜色空间
plot_pixels(data, title="Input color space: 16 million possible colors")

# 使用 K-Means 聚类进行颜色量化，random_state (int): 随机种子，保证结果可复现
kmeans = MiniBatchKMeans(n_clusters=16, n_init='auto', random_state=42)
kmeans.fit(data)

# 将每个像素映射到最近的聚类中心
new_colors = kmeans.cluster_centers_[kmeans.predict(data)]

# 可视化聚类后的颜色空间
plot_pixels(data, colors=new_colors, title="Reduced color space: 16 colors")

# 将聚类后的颜色重新变形为原始图像的形状
china_recolored = new_colors.reshape(china.shape)

# 并排显示原始图像和聚类重构后的图像
fig, ax = plt.subplots(1, 2, figsize=(16, 6), subplot_kw=dict(xticks=[], yticks=[]))
fig.subplots_adjust(wspace=0.05)  # 减小子图之间的间距
ax[0].imshow(china)
ax[0].set_title('Original Image', size=16)
ax[1].imshow(china_recolored)
ax[1].set_title('16-color Image', size=16)
plt.show()